import numpy as np
import pickle
from skimage.io import imread
from skimage.feature import canny
from skimage.filters import sobel, gaussian, median
from scipy.ndimage import uniform_filter
import matplotlib.pyplot as plt
from skimage.color import rgb2gray
from sklearn.metrics import accuracy_score
import joblib  # 确保正确导入 joblib
import time  # 用于记录时间


def extract_features(image):
    # 提取多个特征，包括滤波操作
    features = [
        image.reshape(-1, 1)  # 原始像素值
    ]

    # 加入均值滤波
    mean_filtered = uniform_filter(image, size=3)
    features.append(mean_filtered.reshape(-1, 1))

    # 加入高斯滤波
    gaussian_filtered = gaussian(image, sigma=1)
    features.append(gaussian_filtered.reshape(-1, 1))

    # 加入Sobel滤波
    sobel_filtered = sobel(image)
    features.append(sobel_filtered.reshape(-1, 1))

    # 加入Canny边缘检测
    canny_edges = canny(image)
    features.append(canny_edges.reshape(-1, 1))

    # 打印每个步骤的特征数量
    print(f"每个通道提取后的特征数量: {len(features)}")
    print(f"每个通道处理后的特征维度: {np.concatenate(features, axis=1).shape}")

    # 确保拼接后每个像素有15个特征
    return np.concatenate(features, axis=1)  # 合并所有特征，形成15维的特征向量



def segment():
    # 记录开始时间
    start_time = time.time()

    # 加载图像和模型
    image = imread('Sandstone_2.tif')
    segment = imread('Sandstone_2_segment.tif')

    # 确保图像是RGB格式（如果是灰度图，转换为RGB）
    if image.ndim == 2:  # 如果是灰度图
        image = np.repeat(image[:, :, np.newaxis], 3, axis=2)  # 将灰度图扩展为RGB

    # 获取图像尺寸
    print(f"图像形状: {image.shape}")

    # 提取RGB通道
    R = image[:, :, 0]  # 红色通道
    G = image[:, :, 1]  # 绿色通道
    B = image[:, :, 2]  # 蓝色通道

    # 提取每个通道的特征
    X_r = extract_features(R)  # 对红色通道提取特征
    X_g = extract_features(G)  # 对绿色通道提取特征
    X_b = extract_features(B)  # 对蓝色通道提取特征

    # 合并RGB通道的特征，确保每个像素点都有15个特征
    X = np.concatenate([X_r, X_g, X_b], axis=1)  # 此时X应该有 15 * 像素数 个特征
    print(f"合并后的特征维度: {X.shape}")
   
    # 加载预训练模型
    with open('clf.joblib', 'rb') as f:
        clf = joblib.load(f)  # 使用 joblib.load 来加载模型
        print(f"成功加载模型：{clf}")
    # 记录预测时间
    prediction_start_time = time.time()
    y_pred = clf.predict(X)  # 对所有特征进行预测

    if segment.ndim == 3:  # 如果是RGB图像
        print("标签图像是RGB图像，正在转换为灰度图像...")
        segment = rgb2gray(segment)  # 转换为灰度图像
        segment = (segment * 255).astype(np.uint8)  # 转换为8位灰度图
    
    y = segment.flatten()

    print(f"预测结果 y_pred 维度: {y_pred.shape}")
    print(f"标签 segment 维度: {segment.flatten().shape}")

    # 确保预测结果的形状与标签图像一致
    predicted_segment = y_pred.reshape(segment.shape[:2])  # 重新调整预测结果为图像形状
    prediction_time = time.time() - prediction_start_time

    # 计算准确率
    accuracy = accuracy_score(segment.flatten(), y_pred)

    # 打印分区信息
    print(f"分区准确率: {accuracy}")
    print(f"预测耗时: {prediction_time:.2f}秒")

    # 显示结果
    plt.figure(figsize=(15, 5))

    plt.subplot(1, 3, 1)
    plt.title("Original Image")
    plt.imshow(image, cmap='gray')  # 对于灰度图设置cmap='gray'
    plt.axis('off')

    plt.subplot(1, 3, 2)
    plt.title("Segment")
    plt.imshow(segment, cmap='gray')
    plt.axis('off')

    plt.subplot(1, 3, 3)
    plt.title(f"Segmentation (acc: {accuracy:.3f})")
    plt.imshow(predicted_segment, cmap='gray')
    plt.axis('off')

    plt.tight_layout()
    plt.show()




if __name__ == "__main__":
    segment()
